#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
from PySide6.QtWidgets import QApplication, QMainWindow, QMessageBox
from disable_nagle_ui import Ui_Form
import ctypes
import winreg

YES = QMessageBox.StandardButton.Yes
NO = QMessageBox.StandardButton.No

def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False


def set_nagle(state):
    root_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r'SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces', 0, winreg.KEY_READ)
    num_sub_keys = winreg.QueryInfoKey(root_key)[0]
    for i in range(num_sub_keys):
        skey_name = winreg.EnumKey(root_key, i)
        skey = winreg.OpenKeyEx(root_key, skey_name, 0, winreg.KEY_ALL_ACCESS)
        try:
            winreg.QueryValueEx(skey, 'DhcpIPAddress')
        except FileNotFoundError:
            continue
        if state:
            winreg.SetValueEx(skey, 'TcpAckFrequency', 0, winreg.REG_DWORD, 0x1)
            winreg.SetValueEx(skey, 'TCPNoDelay', 0, winreg.REG_DWORD, 0x1)
        else:
            winreg.SetValueEx(skey, 'TcpAckFrequency', 0, winreg.REG_DWORD, 0x0)
            winreg.SetValueEx(skey, 'TCPNoDelay', 0, winreg.REG_DWORD, 0x0)

    try:
        root2_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r'SOFTWARE\Microsoft\MSMQ\Parameters', 0,
                                   winreg.KEY_ALL_ACCESS)
    except:
        winreg.CreateKeyEx(winreg.HKEY_LOCAL_MACHINE, r'SOFTWARE\Microsoft\MSMQ\Parameters', 0,
                           winreg.KEY_ALL_ACCESS)
        root2_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r'SOFTWARE\Microsoft\MSMQ\Parameters', 0,
                                   winreg.KEY_ALL_ACCESS)
    if state:
        winreg.SetValueEx(root2_key, 'TCPNoDelay', 0, winreg.REG_DWORD, 0x1)
    else:
        winreg.SetValueEx(root2_key, 'TCPNoDelay', 0, winreg.REG_DWORD, 0x0)

    root3_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r'SYSTEM\CurrentControlSet\Services\Tcpip\Parameters', 0, winreg.KEY_ALL_ACCESS)
    if state:
        winreg.SetValueEx(root3_key, 'GlobalMaxTcpWindowSize', 0, winreg.REG_DWORD, 256960)
        winreg.SetValueEx(root3_key, 'DefaultTTL', 0, winreg.REG_DWORD, 0x80)
        winreg.SetValueEx(root3_key, 'EnablePMTUBHDetect', 0, winreg.REG_DWORD, 0)
    else:
        try:
            winreg.DeleteValue(root3_key, 'GlobalMaxTcpWindowSize')
        except:
            pass
        winreg.SetValueEx(root3_key, 'DefaultTTL', 0, winreg.REG_DWORD, 0x40)
        winreg.SetValueEx(root3_key, 'EnablePMTUBHDetect', 0, winreg.REG_DWORD, 1)



def check_and_apply_admin(window_obj):
    if not is_admin():
        QMessageBox.warning(window_obj, '无管理员权限', '申请以管理员权限执行程序', YES, YES)
        ctypes.windll.shell32.ShellExecuteW(None, "runas", sys.executable, __file__, None, 1)
        sys.exit(1)


class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        QMainWindow.__init__(self, parent)
        self.ui = Ui_Form()
        self.ui.setupUi(self)
        self.setFixedSize(self.size())
        self.setWindowTitle('一键禁用Nagle')
        check_and_apply_admin(self)
        self.state = False
        self.update_label_state()
        self.ui.button_switch.clicked.connect(self.slot_button_switch_clicked)

    def slot_button_switch_clicked(self):
        self.state = not self.state
        self.update_label_state()

    def update_label_state(self):
        self.ui.label_state.setText('已禁用Nagle' if self.state else '未禁用Nagle')
        set_nagle(self.state)

    def closeEvent(self, event):
        super().closeEvent(event)
        sel = QMessageBox.question(self, '请选择', '是否立即重启？', YES|NO, NO)
        if sel == YES:
            os.system("shutdown /r /t 1")


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
